// $Id: wfs-loop.c 78 2012-07-20 21:30:18Z alan.watson.f $

#define _POSIX_C_SOURCE 200112L

#include <math.h>

#include "oaxaca-config.h"
#include "oaxaca.h"

#include "wfits.h"
#include "wfs.h"

////////////////////////////////////////////////////////////////////////////////

static mat_t *exposure_z = NULL;

static void
exposure_start(void)
{
  oaxaca_info("exposing.");
  
  int ny = oaxaca_shm->wfs.window.ny;
  int nx = oaxaca_shm->wfs.window.nx;

  if (exposure_z == NULL) {
    exposure_z = mat_alloc(ny, nx);
  } else if (mat_nx(exposure_z) != ny || mat_ny(exposure_z) != ny) {
    mat_free(exposure_z);
    exposure_z = mat_alloc(ny, nx);
  }
  mat_set(exposure_z, 0.0);
}

static void
exposure_next(mat_t *z)
{
  mat_add(exposure_z, z);
}

static void
exposure_end(void)
{
  oaxaca_info("writing exposure to FITS file.");

  FILE *fp = fopen("exposure.fits", "wb");
  if (fp == NULL)
    oaxaca_error("unable to open \"exposure.fits\".");

  oaxaca_wfits_write_basic_header(fp, exposure_z);
  oaxaca_wfits_write_date(fp, "DATE");
  oaxaca_wfits_write_string(fp, "INSTRUME", "Oaxaca WFS");
  oaxaca_wfits_write_string(fp, "DETECTOR", "IDS uEye UI-5540SE-M S/N 4002849489");
  double dit  = 1.0 / oaxaca_shm->wfs.framefrequency;
  oaxaca_wfits_write_double(fp, "EXPTIME", dit, "%.6f");
  oaxaca_wfits_write_double(fp, "DIT", dit, "%.6f");
  double ndit = oaxaca_shm->wfs.cycle;
  oaxaca_wfits_write_double(fp, "NDIT", ndit, "%.0f");
  oaxaca_wfits_write_double(fp, "PIXFREQ", oaxaca_shm->wfs.pixelclockfrequency, "%.1f");
  oaxaca_wfits_write_double(fp, "FRMFREQ", oaxaca_shm->wfs.framefrequency, "%.3f");
  oaxaca_wfits_write_long(fp, "WINSX", oaxaca_shm->wfs.window.sx);
  oaxaca_wfits_write_long(fp, "WINSY", oaxaca_shm->wfs.window.sy);
  oaxaca_wfits_write_long(fp, "WINNX", oaxaca_shm->wfs.window.nx);
  oaxaca_wfits_write_long(fp, "WINNY", oaxaca_shm->wfs.window.ny);

  double frames = oaxaca_shm->wfs.cycle;
  mat_scale(exposure_z, 1.0 / frames);

  // The WFS detector has a 14-bit ADC, so raw pixel values are in the
  // range 0 to 16383. However, we will average images and also may
  // remove darks. Therefore, we select BSCALE = 0.5 and BZERO = 0.0 to
  // allow a BITPIX = 16 FITS file to represent -16384.0 to +16383.5
  // with a resolution of 0.5.
  double bscale = 0.5;
  double bzero = 0.0;

  oaxaca_wfits_write_data(fp, exposure_z, bzero, bscale);

  fclose(fp);

  oaxaca_info("finished exposing.");
}

////////////////////////////////////////////////////////////////////////////////

static mat_t *dark_z = NULL;

static void
dark_start(void)
{
  if (dark_z != NULL)
    mat_free(dark_z);
  dark_z = NULL;
}

static void
dark_next(mat_t *z)
{
  if (dark_z != NULL)
    mat_sub(z, dark_z);
}

static void
dark_end(void)
{
  dark_z = exposure_z;
  exposure_z = NULL;
}

////////////////////////////////////////////////////////////////////////////////

static void
determine_spot_centroid(mat_t *z, oaxaca_spot_t *spot, double *yc, double *xc)
{
  size_t iyc0 = (int) spot->yc0;
  size_t ixc0 = (int) spot->xc0;
  size_t ms = spot->ms;
  size_t mc = spot->mc;

  size_t sys = iyc0 - ms;
  size_t sxs = ixc0 - ms;
  size_t eys = iyc0 + ms;
  size_t exs = ixc0 + ms;
  
  size_t iymax = sys;
  size_t ixmax = sxs;
  double zmax = mat_elt(z, iymax, ixmax);
  for (size_t iy = sys; iy < eys; ++iy)
    for (size_t ix = sxs; ix < exs; ++ix)
      if (mat_elt(z, iy, ix) > zmax) {
        iymax = iy;
        ixmax = ix;
        zmax = mat_elt(z, iy, ix);
      }

  // Make sure the centroid region is within the spot region.
  
  size_t iyc;
  if (iymax < sys + mc)
    iyc = sys + mc;
  else if (iymax > eys - mc)
    iyc = eys - mc;
  else
    iyc = iymax;

  size_t ixc;
  if (ixmax < sxs + mc)
    ixc = sxs + mc;
  else if (ixmax > exs - mc)
    ixc = exs - mc;
  else
    ixc = ixmax;
  
  size_t syc = iyc - mc;
  size_t sxc = ixc - mc;
  size_t eyc = iyc + mc;
  size_t exc = ixc + mc;
  
  // Determine centroids in centroid region around spot.
  double w = 0.0, wx = 0.0, wy = 0.0;
  for (size_t iy = syc; iy <= eyc; ++iy)
    for (size_t ix = sxc; ix <= exc; ++ix) {
      w  += mat_elt(z, iy, ix);
      wy += mat_elt(z, iy, ix) * iy;
      wx += mat_elt(z, iy, ix) * ix;
    }
    
  *yc = wy / w;  
  *xc = wx / w;
}

void
determine_error(mat_t *z, double *mean_dxc, double *mean_dyc)
{

  // Spots for (wfs-set-window! (732 -2 88) 640 (* 2 288) 288)
  double x0 = 407.99 - 1.0 - 121.68;
  double y0 = 118.73 - 1.0 - 113.25;
  oaxaca_spot_t spot[] = {
    { .yc0 = y0 +  28.89, .xc0 = x0 + 119.39, .ms = 20, .mc = 10 },
    { .yc0 = y0 +  26.91, .xc0 = x0 + 169.52, .ms = 20, .mc = 10 },
    { .yc0 = y0 +  70.50, .xc0 = x0 +  73.55, .ms = 20, .mc = 10 },
    { .yc0 = y0 +  69.93, .xc0 = x0 + 119.93, .ms = 20, .mc = 10 },
    { .yc0 = y0 +  68.97, .xc0 = x0 + 168.06, .ms = 20, .mc = 10 },
    { .yc0 = y0 +  66.69, .xc0 = x0 + 213.51, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 113.45, .xc0 = x0 +  29.78, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 113.63, .xc0 = x0 +  75.31, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 113.25, .xc0 = x0 + 121.68, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 113.41, .xc0 = x0 + 168.79, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 112.52, .xc0 = x0 + 213.32, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 110.77, .xc0 = x0 + 256.87, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 158.32, .xc0 = x0 +  31.29, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 158.68, .xc0 = x0 +  76.42, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 158.68, .xc0 = x0 + 123.51, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 159.93, .xc0 = x0 + 170.66, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 160.08, .xc0 = x0 + 214.28, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 159.23, .xc0 = x0 + 256.58, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 201.21, .xc0 = x0 +  76.89, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 201.82, .xc0 = x0 + 128.40, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 206.91, .xc0 = x0 + 173.30, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 208.47, .xc0 = x0 + 216.29, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 253.46, .xc0 = x0 + 133.67, .ms = 20, .mc = 10 },
    { .yc0 = y0 + 257.15, .xc0 = x0 + 178.52, .ms = 20, .mc = 10 },
  };
  size_t nspot = sizeof(spot) / sizeof(*spot);
  
  double dyc[nspot];
  double dxc[nspot];
  
  //oaxaca_info("%lu", (unsigned long) oaxaca_shm->wfs.cycle);
  for (int i = 0; i < nspot; ++i) {
    double xc, yc;
    determine_spot_centroid(z, &spot[i], &yc, &xc);
    dyc[i] = yc - spot[i].yc0;
    dxc[i] = xc - spot[i].xc0;
    //oaxaca_info("%2d %5.1f %5.1f %+5.1f %+5.1f ", i, yc, xc, dyc[i], dxc[i]);
  }
  
  *mean_dyc = 0.0;
  *mean_dxc = 0.0;
  for (int i = 0; i < nspot; ++i) {
    *mean_dyc += dyc[i];
    *mean_dxc += dxc[i];
  }
  *mean_dyc /= (double) nspot;
  *mean_dxc /= (double) nspot;

  //oaxaca_info("%lu %+6.2f %+6.2f", oaxaca_shm->wfs.cycle, *mean_dxc, *mean_dyc);
}

////////////////////////////////////////////////////////////////////////////////

static void
correct(double mean_dxc, double mean_dyc)
{
  if (oaxaca_shm->wfc.command == 'I') {
    double aogain = oaxaca_shm->wfs.aogain;
    if (isfinite(mean_dxc))
      oaxaca_shm->wfc.mode[0] += aogain * mean_dxc;
    if (isfinite(mean_dyc))
      oaxaca_shm->wfc.mode[1] += aogain * mean_dyc;
    if (oaxaca_shm->wfs.cycle % 40 == 0) {
      oaxaca_shm->wfc.mode[0] = 4.0 * ((double) rand() / (double) RAND_MAX - 0.5);
      oaxaca_shm->wfc.mode[1] = 4.0 * ((double) rand() / (double) RAND_MAX - 0.5);
    }
    oaxaca_shm->wfc.modal = 1;
    oaxaca_shm->wfc.command = 't';
  }
}
  
////////////////////////////////////////////////////////////////////////////////

static unsigned long exposureframes;

int
loop_frame_handler(mat_t *z)
{
  if (oaxaca_shm->wfs.command == 'i') {
    oaxaca_info("aborting loop.");
    return 0;
  }

  ++oaxaca_shm->wfs.cycle;
  
  dark_next(z);

  if (oaxaca_shm->wfs.command == 'A') {
    double mean_dxc, mean_dyc;
    determine_error(z, &mean_dxc, &mean_dyc);
    correct(mean_dxc, mean_dyc);
  }
  
  exposure_next(z);
  if (oaxaca_shm->wfs.cycle % exposureframes == 0) {
    exposure_end();
    if (oaxaca_shm->wfs.command != 'A')
      return 0;
    exposure_start();
  }

  return 1;
}

void
oaxaca_wfs_loop(void)
{
  oaxaca_info("looping.");

  oaxaca_shm->wfs.cycle = 0;

  if (oaxaca_shm->wfs.command == 'D')
    dark_start();
  exposure_start();
  exposureframes = (unsigned long)
    ceil(oaxaca_shm->wfs.exposuretime * oaxaca_shm->wfs.framefrequency);

  oaxaca_ids_run(loop_frame_handler);

  if (oaxaca_shm->wfs.command == 'D')
    dark_end();

  oaxaca_info("finished looping.");
}

////////////////////////////////////////////////////////////////////////////////
